# !usr/bin/env python
# -*- coding:utf-8 -*-

'''
 Description  : 
 Version      : 1.0
 Author        : qww
 Date         : 2021-05-02 16:57:13
 LastEditors  : qww
 LastEditTime : 2021-05-02 19:32:00
'''

import os
from PIL import Image
import torch
from torchvision import transforms


CLASS_NAMES = ['ants', 'bees']


def predict(filename: str, model_path: str) -> str:
    img = Image.open(fp=filename)
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    model = torch.load(model_path)
    model = model.cpu()
    with torch.no_grad():
        output = model(transform(img).view(-1, 3, 224, 224))
        _, pred = torch.max(output, 1)

    os.remove(filename)
    return CLASS_NAMES[pred.item()]
